from torchvision.datasets import MNIST
from torchvision.transforms import Compose, ToTensor, Normalize, ToPILImage
from torch.utils.data import DataLoader
import torch
import torchvision
from torch import nn
import torch.nn.functional as F
from torch.optim import Adam
import os
import numpy as np
import matplotlib.pyplot as plt
import cv2

# TODO:1.转换成GPU形式

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# TODO:2.修改网络模型


BATCH_SIZE = 128
TEST_BATCH_SIZE = 1000


# 1、准备数据集
def get_dataloader(train=True, batch_size=BATCH_SIZE):
    transform_fn = Compose([
        ToTensor(),
        Normalize(mean=(0.1307,), std=(0.3081,))
    ])
    dataset = MNIST(root='./data', train=train, download=False, transform=transform_fn)
    data_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, num_workers=2)
    return data_loader


# 2、搭建网络
# 平均准确率， 平均损失 0.9585 0.14125641956925392
class mnistNet(nn.Module):
    def __init__(self):
        super(mnistNet, self).__init__()
        self.fc1 = nn.Linear(1 * 28 * 28, 28)
        self.fc2 = nn.Linear(28, 10)

    def forward(self, x):
        x = x.view(-1, 1 * 28 * 28)
        x = self.fc1(x)
        x = F.relu(x)
        out = self.fc2(x)
        return out


# 修改网络
# 平均准确率， 平均损失 0.9917000472545624 0.028591989167034627
class mnistNet2(nn.Module):
    def __init__(self):
        super(mnistNet2, self).__init__()
        # 输入通道1，输出通道16，卷积核5，步长1， padding2
        self.conv1 = nn.Sequential(  # 1x28x28
            nn.Conv2d(
                in_channels=1,
                out_channels=16,
                kernel_size=5,
                stride=1,
                padding=2,
            ),  # 16x28x28 ==》16 * [(28-5+4+1)/1] * [(28-5+4+1)/1]
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)  # 16*14*14
        )
        self.conv2 = nn.Sequential(  # 16x14x14
            nn.Conv2d(16, 32, 5, 1, 2),  # 32x14x14
            nn.ReLU(),
            nn.MaxPool2d(2),  # 32x7x7
        )
        self.out = nn.Linear(32 * 7 * 7, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)  # 展平多维的卷积图成 (batch_size, 32 * 7 * 7)
        output = self.out(x)
        return output


# 实例化模型，实例化优化器，实例化损失函数
mnist_net = mnistNet2().to(device)
optimizer = Adam(mnist_net.parameters(), lr=0.001)
loss_func = nn.CrossEntropyLoss()


# 定义训练器
def train(epoches: int):
    '''
    训练网络
    :param epoches:
    :return:
    '''
    train_loss_list = []

    for epoch in range(epoches):
        train_dataloader = get_dataloader()
        epoch_loss = 0
        for idx, (img, label) in enumerate(train_dataloader):
            img, label = img.to(device), label.to(device)
            output = mnist_net(img)
            loss = loss_func(output, label)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if idx % 100 == 0:
                torch.save(mnist_net.state_dict(), 'model/model2.pkl')
                torch.save(optimizer.state_dict(), 'model/optimizer2.pkl')
                print('Train Epoch:{} [{}/{} ({:.0f}%)]\tLoss:{:.6f}'.format(
                    epoch, idx * len(img), len(train_dataloader.dataset),
                           100. * idx / len(train_dataloader), loss.cpu().item()))
                epoch_loss += loss.cpu().item()
        train_loss_list.append(np.mean(epoch_loss))
    show_loss(train_loss_list)


def test():
    mnist_net.eval()
    loss_list = []
    acc_list = []
    test_dataloader = get_dataloader(train=False, batch_size=TEST_BATCH_SIZE)
    for idx, (img, label) in enumerate(test_dataloader):
        img, label = img.to(device), label.to(device)
        with torch.no_grad():
            output = mnist_net(img)
            cur_loss = loss_func(output, label)
            loss_list.append(cur_loss.cpu().item())
            pred = output.max(dim=-1)[-1]
            cur_acc = pred.eq(label).float().mean()
            acc_list.append(cur_acc.cpu().item())
    print("平均准确率， 平均损失", np.mean(acc_list), np.mean(loss_list))


# 实现单张图片的可视化
def show_pic(images, labels):
    img = torchvision.utils.make_grid(images)
    img = img.numpy().transpose(1, 2, 0)
    std = [0.5, 0.5, 0.5]
    mean = [0.5, 0.5, 0.5]
    img = img * std + mean
    print(labels)
    cv2.imshow('win', img)
    key_pressed = cv2.waitKey(0)


# 预测一张图片
def predict_onepic():
    mnist_net.eval()
    test_dataloader = get_dataloader(train=False, batch_size=1)
    image, label = next(iter(test_dataloader))
    print(image.size())
    show_pic(image, label)  # 展示图片
    img, label = image.to(device), label.to(device)
    with torch.no_grad():
        output = mnist_net(img)
        # cur_loss = loss_func(output, label)
        pred = output.max(dim=-1)[-1]
        print('predict:', pred)


# 画图，loss随时间的变化图
def show_loss(loss_list):
    plt.figure(figsize=(20, 8))
    plt.plot(range(len(loss_list)), loss_list)
    plt.show()


if __name__ == '__main__':
    if os.path.exists('model/model2.pkl'):
        mnist_net.load_state_dict(torch.load('model/model2.pkl'))
        optimizer.load_state_dict((torch.load('model/optimizer2.pkl')))
    # train(10)
    # test()
    predict_onepic()
